{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem setting\n",
    "\n",
    "In this tutorial, we demonstrate how graph neural networks can be used for recommendation. Here we focus on item-based recommendation model. This method in this tutorial recommends items that are similar to the ones purchased by the user. We demonstrate the recommendation model on the MovieLens dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get started\n",
    "\n",
    "DGL can be used with different deep learning frameworks. Currently, DGL can be used with Pytorch and MXNet. Here, we show how DGL works with Pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we load DGL, we need to set the DGL backend for one of the deep learning frameworks. Because this tutorial develops models in Pytorch, we have to set the DGL backend to Pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "from dgl import DGLGraph\n",
    "\n",
    "# Load Pytorch as backend\n",
    "dgl.load_backend('pytorch')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the rest of necessary libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy import stats\n",
    "from scipy import sparse as spsp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset\n",
    "\n",
    "We use the MoiveLens dataset for demonstration because it is commonly used for recommendation models. In this dataset, there are two types of nodes: users and movies. The movie nodes have three attributes: year, title and genre. There are ratings between user nodes and movie nodes. Each rating has a timestamp. In our recommendation model, we don't consider ratings and timestamps.\n",
    "\n",
    "**Note**: It is not necessarily the best dataset to demonstrate the power of GNN for recommendation. We have prepared the dataset to simplify the demonstration.\n",
    "\n",
    "To run the data preprocessing script, a user needs to download the English dictionary of the stanfordnlp package first. However, the following command only needs to run once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please uncomment the two commands when the tutorial is run for the first time.\n",
    "#import stanfordnlp\n",
    "#stanfordnlp.download('en')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the MovieLens dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from movielens import MovieLens\n",
    "data = MovieLens('.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate some statistics of the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ratings = data.ratings\n",
    "user_id = np.array(ratings['user_idx'])\n",
    "movie_id = np.array(ratings['movie_idx'])\n",
    "user_movie_spm = spsp.coo_matrix((np.ones((len(user_id),)), (user_id, movie_id)))\n",
    "num_users, num_movies = user_movie_spm.shape\n",
    "print('#user-movie iterations:', len(movie_id))\n",
    "print('#users:', num_users)\n",
    "print('#movies:', num_movies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We split the dataset into a training set and a testing set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ratings_train = ratings[~(ratings['valid_mask'] | ratings['test_mask'])]\n",
    "user_latest_item_indices = (\n",
    "        ratings_train.groupby('user_id')['timestamp'].transform(pd.Series.max) ==\n",
    "        ratings_train['timestamp'])\n",
    "user_latest_item = ratings_train[user_latest_item_indices]\n",
    "user_latest_item = dict(\n",
    "        zip(user_latest_item['user_idx'].values, user_latest_item['movie_idx'].values))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct the training, validation and testing dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The training dataset\n",
    "user_id = np.array(ratings_train['user_idx'])\n",
    "movie_id = np.array(ratings_train['movie_idx'])\n",
    "user_movie_spm = spsp.coo_matrix((np.ones((len(user_id),)), (user_id, movie_id)))\n",
    "assert num_users == user_movie_spm.shape[0]\n",
    "assert num_movies == user_movie_spm.shape[1]\n",
    "train_size = len(user_id)\n",
    "print('#training size:', train_size)\n",
    "\n",
    "# The validation and testing dataset\n",
    "users_valid = ratings[ratings['valid_mask']]['user_idx'].values\n",
    "movies_valid = ratings[ratings['valid_mask']]['movie_idx'].values\n",
    "users_test = ratings[ratings['test_mask']]['user_idx'].values\n",
    "movies_test = ratings[ratings['test_mask']]['movie_idx'].values\n",
    "valid_size = len(users_valid)\n",
    "test_size = len(users_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The recommendation model\n",
    "\n",
    "At large, the model first learns item embeddings from the user-item interaction dataset and use the item embeddings to recommend users similar items they have purchased. To learn item embeddings, we first need to construct an item similarity graph and train GNN on the item graph.\n",
    "\n",
    "There are many ways of constructing the item similarity graph. Here we use the [SLIM model](https://dl.acm.org/citation.cfm?id=2118303) to learn item similarity and use the learned result to construct the item graph. The resulting graph will have an edge between two items if they are similar and the edge has a weight that represents the similarity score.\n",
    "\n",
    "After the item similarity graph is constructed, we run a GNN model on it and use the vertex connectivity as the training signal to train the GNN model. The GNN training procedure is very similar to the link prediction task in [the previous section](https://github.com/zheng-da/DGL_devday_tutorial/blob/master/BasicTasks_pytorch.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct the movie graph with SLIM\n",
    "SLIM is an item-based recommendation model. When training SLIM on a user-item dataset, it learns an item similarity graph. This similarity graph is the item graph we construct for the GNN model.\n",
    "\n",
    "Please follow the instruction on the [SLIM github repo](https://github.com/KarypisLab/SLIM) to install SLIM.\n",
    "\n",
    "The SLIM only needs to run once and can be saved on the disk for future experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from SLIM import SLIM, SLIMatrix\n",
    "model = SLIM()\n",
    "params = {'algo': 'cd', 'nthreads': 2, 'l1r': 2.0, 'l2r': 1.0}\n",
    "trainmat = SLIMatrix(user_movie_spm.tocsr())\n",
    "model.train(params, trainmat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the SLIM similarity matrix into DGL. We store the vertex similarity as edge data on DGL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_spm = model.to_csr()\n",
    "print('#edges:', movie_spm.nnz)\n",
    "print('most similar:', np.max(movie_spm.data))\n",
    "print('most unsimilar:', np.min(movie_spm.data))\n",
    "\n",
    "g = dgl.DGLGraph(movie_spm, readonly=True)\n",
    "g.edata['similarity'] = torch.tensor(movie_spm.data, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct the item features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "year = np.expand_dims(data.movie_data['year'], axis=1)\n",
    "genre = data.movie_data['genre']\n",
    "title = data.movie_data['title']\n",
    "features = torch.tensor(np.concatenate((genre, title), axis=1), dtype=torch.float32)\n",
    "print('#features:', features.shape[1])\n",
    "in_feats = features.shape[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GNN models\n",
    "\n",
    "We run GNN on the item graph to compute item embeddings. In this tutorial, we use a customized [GraphSage](https://cs.stanford.edu/people/jure/pubs/graphsage-nips17.pdf) model to compute node embeddings. The original GraphSage performs the following computation on every node $v$ in the graph:\n",
    "\n",
    "$$h_{N(v)}^{(l)} \\gets AGGREGATE_k({h_u^{(l-1)}, \\forall u \\in N(v)})$$\n",
    "$$h_v^{(l)} \\gets \\sigma(W^k \\cdot CONCAT(h_v^{(l-1)}, h_{N(v)}^{(l)})),$$\n",
    "\n",
    "where $N(v)$ is the neighborhood of node $v$ and $l$ is the layer Id.\n",
    "\n",
    "The original GraphSage model treats each neighbor equally. However, the SLIM model learns the item similarity based on the user-item iteration. The GNN model should take the similarity into account. Thus, we customize the GraphSage model in the following fashion. Instead of aggregating all neighbors equally, we aggregate neighbors embeddings rescaled by the similarity on the edges. Thus, the aggregation step is defined as follows:\n",
    "\n",
    "$$h_{N(v)}^{(l)} \\gets \\Sigma_{u \\in N(v)}({h_u^{(l-1)} * s_{uv}}),$$\n",
    "\n",
    "where $s_{uv}$ is the similarity score between two vertices $u$ and $v$.\n",
    "\n",
    "The GNN model has multiple layers. In each layer, a vertex accesses its direct neighbors. When we stack $k$ layers in a model, a node $v$ access neighbors within $k$ hops. The output of the GNN model is node embeddings that represent the nodes and all information in the k-hop neighborhood.\n",
    "\n",
    "<img src=\"https://github.com/zheng-da/DGL_devday_tutorial/raw/master/GNN.png\" alt=\"drawing\" width=\"600\"/>\n",
    "\n",
    "We implement the computation in each layer of the customized GraphSage model in `SAGEConv` and implement the multi-layer model in `GraphSAGEModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sageconv import SAGEConv\n",
    "\n",
    "class GraphSAGEModel(nn.Module):\n",
    "    def __init__(self,\n",
    "                 in_feats,\n",
    "                 n_hidden,\n",
    "                 out_dim,\n",
    "                 n_layers,\n",
    "                 activation,\n",
    "                 dropout,\n",
    "                 aggregator_type):\n",
    "        super(GraphSAGEModel, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "\n",
    "        # input layer\n",
    "        self.layers.append(SAGEConv(in_feats, n_hidden, aggregator_type,\n",
    "                                    feat_drop=dropout, activation=activation))\n",
    "        # hidden layers\n",
    "        for i in range(n_layers - 1):\n",
    "            self.layers.append(SAGEConv(n_hidden, n_hidden, aggregator_type,\n",
    "                                        feat_drop=dropout, activation=activation))\n",
    "        # output layer\n",
    "        self.layers.append(SAGEConv(n_hidden, out_dim, aggregator_type,\n",
    "                                    feat_drop=dropout, activation=None))\n",
    "\n",
    "    def forward(self, g, features):\n",
    "        h = features\n",
    "        for layer in self.layers:\n",
    "            h = layer(g, h, g.edata['similarity'])\n",
    "        return h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Item Embeddings\n",
    "\n",
    "We train the item embeddings with the edges in the item graph as the training signal. This step is very similar to the link prediction task in the [basic applications](https://github.com/zheng-da/DGL_devday_tutorial/blob/master/BasicTasks_pytorch.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because the MovieLens dataset has sparse features (both genre and title are stored as multi-hot encoding). The sparse features have many dimensions. To run GNN on the item features, we first create an encoding layer to project the sparse features to a lower dimension. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncodeLayer(nn.Module):\n",
    "    def __init__(self, in_feats, num_hidden):\n",
    "        super(EncodeLayer, self).__init__()\n",
    "        self.proj = nn.Linear(in_feats, num_hidden)\n",
    "        \n",
    "    def forward(self, feats):\n",
    "        return self.proj(feats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We simply use node connectivity as the training signal: nodes connected by edges are similar, while nodes not connected by edges are dissimilar.\n",
    "\n",
    "To train such a model, we need to deploy negative sampling to construct negative samples. A positive sample is an edge that exist in the original graph, while a negative sample is a pair of nodes that don't have an edge between them in the graph. We usually train on each positive sample with multiple negative samples.\n",
    "\n",
    "After having the node embeddings, we compute the similarity scores on positive samples and negative samples. We construct the following loss function on a positive sample and the corresponding negative samples:\n",
    "\n",
    "$$L = -log(\\sigma(z_u^T z_v)) - Q \\cdot E_{v_n \\sim P_n(v)}(log(\\sigma(-z_u^T z_{v_n}))),$$\n",
    "\n",
    "where $Q$ is the number of negative samples.\n",
    "\n",
    "With this loss, training should increase the similarity scores on the positive samples and decrease the similarity scores on negative samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 0\n",
    "gamma = 0\n",
    "\n",
    "class FISM(nn.Module):\n",
    "    def __init__(self, user_movie_spm, gconv_p, gconv_q, in_feats, num_hidden):\n",
    "        super(FISM, self).__init__()\n",
    "        self.encode = EncodeLayer(in_feats, num_hidden)\n",
    "        self.num_users = user_movie_spm.shape[0]\n",
    "        self.num_movies = user_movie_spm.shape[1]\n",
    "        self.b_u = nn.Parameter(torch.zeros(num_users))\n",
    "        self.b_i = nn.Parameter(torch.zeros(num_movies))\n",
    "        self.user_deg = torch.tensor(user_movie_spm.dot(np.ones(num_movies)))\n",
    "        values = user_movie_spm.data\n",
    "        indices = np.vstack((user_movie_spm.row, user_movie_spm.col))\n",
    "        indices = torch.LongTensor(indices)\n",
    "        values = torch.FloatTensor(values)\n",
    "        self.user_item_spm = torch.sparse_coo_tensor(indices, values, user_movie_spm.shape)\n",
    "        self.users = user_movie_spm.row\n",
    "        self.movies = user_movie_spm.col\n",
    "        self.ratings = user_movie_spm.data\n",
    "        self.gconv_p = gconv_p\n",
    "        self.gconv_q = gconv_q\n",
    "\n",
    "    def _est_rating(self, P, Q, user_idx, item_idx):\n",
    "        bu = self.b_u[user_idx]\n",
    "        bi = self.b_i[item_idx]\n",
    "        user_emb = torch.sparse.mm(self.user_item_spm, P)\n",
    "        user_emb = user_emb[user_idx] / torch.unsqueeze(self.user_deg[user_idx], 1)\n",
    "        tmp = torch.mul(user_emb, Q[item_idx])\n",
    "        r_ui = bu + bi + torch.sum(tmp, 1)\n",
    "        return r_ui\n",
    "    \n",
    "    def est_rating(self, g, features, user_idx, item_idx, neg_item_idx):\n",
    "        h = self.encode(features)\n",
    "        P = self.gconv_p(g, h)\n",
    "        Q = self.gconv_q(g, h)\n",
    "        r = self._est_rating(P, Q, user_idx, item_idx)\n",
    "        neg_sample_size = len(neg_item_idx) / len(user_idx)\n",
    "        neg_r = self._est_rating(P, Q, np.repeat(user_idx, neg_sample_size), neg_item_idx)\n",
    "        return torch.unsqueeze(r, 1), neg_r.reshape((-1, int(neg_sample_size)))\n",
    "\n",
    "    def loss(self, P, Q, r_ui, neg_r_ui):\n",
    "        diff = 1 - (r_ui - neg_r_ui)\n",
    "        return torch.sum(torch.mul(diff, diff)/2) \\\n",
    "            + beta/2 * torch.sum(torch.mul(P, P) + torch.mul(Q, Q)) \\\n",
    "            + gamma/2 * (torch.sum(torch.mul(self.b_u, self.b_u)) + torch.sum(torch.mul(self.b_i, self.b_i)))\n",
    "\n",
    "    def forward(self, g, features, neg_sample_size):\n",
    "        h = self.encode(features)\n",
    "        P = self.gconv_p(g, h)\n",
    "        Q = self.gconv_q(g, h)\n",
    "        tot = len(self.users)\n",
    "        pos_idx = np.random.choice(tot, int(tot/10))\n",
    "        user_idx = self.users[pos_idx]\n",
    "        item_idx = self.movies[pos_idx]\n",
    "        neg_item_idx = np.random.choice(self.num_movies, len(pos_idx) * neg_sample_size)\n",
    "        r_ui = self._est_rating(P, Q, user_idx, item_idx)\n",
    "        neg_r_ui = self._est_rating(P, Q, np.repeat(user_idx, neg_sample_size), neg_item_idx)\n",
    "        r_ui = torch.unsqueeze(r_ui, 1)\n",
    "        neg_r_ui = neg_r_ui.reshape((-1, int(neg_sample_size)))\n",
    "        return self.loss(P, Q, r_ui, neg_r_ui)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We evaluate the performance of the trained item embeddings in the item-based recommendation task. We use the last item that a user purchased to represent the user and compute the similarity between the last item and a list of items (an item the user will purchase and a set of randomly sampled items). We calculate the ranking of the item that will be purchased among the list of items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def RecEvaluate(model, g, features, users_eval, movies_eval, neg_sample_size):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        neg_movies_eval = data.neg_valid[users_eval].flatten()\n",
    "        r, neg_r = model.est_rating(g, features, users_eval, movies_eval, neg_movies_eval)\n",
    "        hits10 = (torch.sum(neg_r > r, 1) <= 10).numpy()\n",
    "        print('HITS@10:{:.4f}'.format(np.mean(hits10)))\n",
    "        return np.mean(hits10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we put everything in the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Model hyperparameters\n",
    "n_hidden = 16\n",
    "n_layers = 1\n",
    "dropout = 0.5\n",
    "aggregator_type = 'sum'\n",
    "\n",
    "# create GraphSAGE model\n",
    "gconv_p = GraphSAGEModel(n_hidden,\n",
    "                         n_hidden,\n",
    "                         n_hidden,\n",
    "                         n_layers,\n",
    "                         F.relu,\n",
    "                         dropout,\n",
    "                         aggregator_type)\n",
    "\n",
    "gconv_q = GraphSAGEModel(n_hidden,\n",
    "                         n_hidden,\n",
    "                         n_hidden,\n",
    "                         n_layers,\n",
    "                         F.relu,\n",
    "                         dropout,\n",
    "                         aggregator_type)\n",
    "\n",
    "model = FISM(user_movie_spm, gconv_p, gconv_q, in_feats, n_hidden)\n",
    "\n",
    "\n",
    "# Training hyperparameters\n",
    "weight_decay = 5e-4\n",
    "n_epochs = 100\n",
    "lr = 1e-3\n",
    "neg_sample_size = 20\n",
    "\n",
    "# use optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "# initialize graph\n",
    "dur = []\n",
    "prev_acc = 0\n",
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "    loss = model(g, features, neg_sample_size)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    print(\"Epoch {:05d} | Loss {:.4f}\".format(epoch, loss.item()))\n",
    "    \n",
    "    acc = RecEvaluate(model, g, features, users_valid, movies_valid, neg_sample_size)\n",
    "\n",
    "print()\n",
    "# Let's save the trained node embeddings.\n",
    "RecEvaluate(model, g, features, users_test, movies_test, neg_sample_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
